import math
import numpy as np # np数学库
import scipy.linalg as la # 线性代数求解库
from pathPreProcess import pathPreProcess
    
class LQRControl:
    def __init__(self, path):
        self.Q = np.eye(5)
        # """较好"""      
        # self.Q[0, 0] = 500
        # self.Q[1, 1] = 50
        # self.Q[2, 2] = 1
        # self.Q[3, 3] = 10
        
        """较好"""      
        self.Q[0, 0] = 1000
        self.Q[1, 1] = 10
        self.Q[2, 2] = 10
        self.Q[3, 3] = 10
        
        self.R = np.eye(2)
        self.dt = 0.1
        self.STOP_FLAG = False
        self.k = 0.1
        self.Lfc = 0.12
        self.L = 0.145
        self.pe = 0
        self.pth_e = 0
        self.i = 0
        self.traj, self.traj_list, self.Point_O = pathPreProcess(path)

    def lqr_speed_steering_control(self, state):
        traj = self.traj_list[self.i]

        # *************** 头 *************************
        self.traj_len = len(traj.x)
        self.ind, e = self.calc_nearest_index(state, traj.x, traj.y, traj.yaw) # e表示车辆离轨迹最近点的距离（也可以理解为位置误差），ind为该点的索引
        tv = traj.v[self.ind]

        k = traj.kappa[self.ind] # kappa为曲率 
        v = state.v
        th_e = self.pi_2_pi(state.yaw - traj.yaw[self.ind])

        # A = [1.0, dt, 0.0, 0.0, 0.0
        #      0.0, 0.0, v, 0.0, 0.0]
        #      0.0, 0.0, 1.0, dt, 0.0]
        #      0.0, 0.0, 0.0, 0.0, 0.0]
        #      0.0, 0.0, 0.0, 0.0, 1.0]
        A = np.zeros((5, 5))
        A[0, 0] = 1.0
        A[0, 1] = self.dt
        A[1, 2] = v
        A[2, 2] = 1.0
        A[2, 3] = self.dt
        A[4, 4] = 1.0

        # B = [0.0, 0.0
        #     0.0, 0.0
        #     0.0, 0.0
        #     v/L, 0.0
        #     0.0, dt]
        B = np.zeros((5, 2))
        B[3, 0] = v / self.L
        B[4, 1] = self.dt

        K, _, _ = self.dlqr(A, B, self.Q, self.R)

        # state vector
        # x = [e, dot_e, th_e, dot_th_e, delta_v]
        # e: lateral distance to the path
        # dot_e: derivative of e
        # th_e: angle difference to the path
        # dot_th_e: derivative of th_e
        # delta_v: difference between current speed and target speed
        x = np.zeros((5, 1))
        x[0, 0] = e
        x[1, 0] = (e - self.pe) / self.dt
        x[2, 0] = th_e
        x[3, 0] = (th_e - self.pth_e) / self.dt
        x[4, 0] = v - tv

        self.pe = e
        self.pth_e - th_e
        # input vector
        # u = [delta, accel]
        # delta: steering angle
        # accel: acceleration
        ustar = -K @ x # @表示矩阵乘法 计算得到控制矩阵u

        # calc steering input 基于跟踪误差模型（考虑道路弯曲）龚建伟MPC/p30
        ff = math.atan2(self.L * k, 1)  # feedforward steering angle  航向误差需要的偏转角
        fb = self.pi_2_pi(ustar[0, 0])  # feedback steering angle   横向误差需要的偏转角
        self.delta = ff + fb
        # print(self.delta)

        if traj.direction[self.ind] < 0:
            self.delta *= -1
        # calc accel input
        self.accel = ustar[1, 0]

        # ***************** 尾 ******************
        if self.ind >= (len(self.traj_list[self.i].x) - 1):
            # 第i段结束
            self.i += 1

        v, w = self.cal_v_w(traj)
        self.is_stop()

        return v, w
    
    def calc_nearest_index(self, state, cx, cy, cyaw):
        dx = [state.x - icx for icx in cx]
        dy = [state.y - icy for icy in cy]

        d = [idx ** 2 + idy ** 2 for (idx, idy) in zip(dx, dy)]

        min_d = min(d)

        ind = d.index(min_d)

        min_d = math.sqrt(min_d)

        dxl = cx[ind] - state.x
        dyl = cy[ind] - state.y

        angle = self.pi_2_pi(cyaw[ind] - math.atan2(dyl, dxl))
        if angle < 0: # 小于0，车辆在轨迹的右侧，大于0，车辆在轨迹的左侧
            min_d *= -1

        return ind, min_d

    def pi_2_pi(self, angle):
        return (angle + math.pi) % (2 * math.pi) - math.pi

    def dlqr(self, A, B, Q, R):
        """Solve the discrete time lqr controller.
        x[k+1] = A x[k] + B u[k]
        cost = sum x[k].T*Q*x[k] + u[k].T*R*u[k]
        # ref Bertsekas, p.151
        """

        # first, try to solve the ricatti equation
        P = self.solve_dare(A, B, Q, R)

        # compute the LQR gain
        K = la.inv(B.T @ P @ B + R) @ (B.T @ P @ A)

        eig_result = la.eig(A - B @ K)

        return K, P, eig_result[0]

    def solve_dare(self, A, B, Q, R):
        """
        solve a discrete time_Algebraic Riccati equation (DARE)
        """
        p = Q
        p_next = Q
        max_iter = 150
        eps = 0.01

        for i in range(max_iter):
            p_next = A.T @ p @ A - A.T @ p @ B @ \
                    la.inv(R + B.T @ p @ B) @ B.T @ p @ A + Q
            if (abs(p_next - p)).max() < eps:
                break
            p = p_next

        return p_next
    
    def cal_v_w(self, traj):
        v = traj.direction[self.ind] * traj.v[self.ind] + self.accel * 0.1
        w = traj.v[self.ind] * math.tan(self.delta) / self.L
        return v, w
    
    def is_stop(self):
        if self.i == len(self.traj_list) and self.ind >= (self.traj_len - 1): # 必须减1
            self.STOP_FLAG = True
